package middleware

import (
	"context"
	"fmt"
	"golang.org/x/time/rate"
	"net"
	"net/http"
	"sync"
	"time"
)

var ipLimiterMux sync.RWMutex

func IpLimiterMiddleware(limiter map[string]*rate.Limiter, frequent time.Duration) func(next http.HandlerFunc) http.HandlerFunc {
	return func(next http.HandlerFunc) http.HandlerFunc {
		return func(w http.ResponseWriter, r *http.Request) {
			ip, _, err := net.SplitHostPort(r.RemoteAddr)
			if err != nil {
				//http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
				msg := fmt.Sprintf(`{"status": 500, "msg": "%s"}`, err.Error())
				fmt.Println("[Error]查询客户端ip地址:", err.Error())
				http.Error(w, msg, http.StatusOK)
				return
			}

			// 从上下文中获取限流器
			ctx := r.Context()
			rl, ok := ctx.Value("IpLimiter").(map[string]*rate.Limiter)
			if !ok {
				// 如果上下文中没有限流器，则初始化
				rl = limiter
				ctx = context.WithValue(ctx, "IpLimiter", rl)
				r = r.WithContext(ctx)
			}

			ipLimiterMux.Lock()
			l, ok := rl[ip]
			if !ok {
				limit := rate.Every(frequent * time.Millisecond) //每100毫秒
				l = rate.NewLimiter(limit, 10)                   //b:桶容量，设置为10个
				rl[ip] = l
			}
			ipLimiterMux.Unlock()

			// 进行限流控制
			if !l.Allow() {
				// 超过限制，返回错误
				//w.WriteHeader(http.StatusOK)
				// StatusTooManyRequests              = 429 // RFC 6585, 4
				//w.Write([]byte(`{"status": 429, "msg": "too many requests"}`))
				fmt.Printf("[Error]ip[%s]限流\n", ip)
				http.Error(w, `{"status": 429, "msg": "too many requests"}`, http.StatusOK)
				return
			}

			// 继续执行下一个中间件或处理函数
			next(w, r)
		}
	}
}
